import cv2
import os
import json
import numpy as np

# 文件夹路径
image_dir = "/media/champrin/Data/Desktop/rune dataset/extract/all2_yolo/set"  # 图片文件夹路径
output_dir = "/media/champrin/Data/Desktop/rune dataset/extract/all2_yolo/set/output"  # 输出图像文件夹路径

os.makedirs(output_dir, exist_ok=True)


drawing = False
last_point = None
mouse_x, mouse_y = -1, -1
circle_radius = 10

def mouse_callback(event, x, y, flags, param):
    global drawing, last_point, mouse_x, mouse_y

    mouse_x, mouse_y = x, y

    if event == cv2.EVENT_LBUTTONDOWN:
        drawing = True
        last_point = (x, y)

    elif event == cv2.EVENT_MOUSEMOVE:
        if drawing and last_point is not None:
            thickness = max(1, int(2 * circle_radius))
            cv2.line(img, last_point, (x, y), (0, 0, 0), thickness)
            last_point = (x, y)

    elif event == cv2.EVENT_LBUTTONUP:
        drawing = False
        last_point = None

cv2.namedWindow("Paint Window")
cv2.setMouseCallback("Paint Window", mouse_callback)

# 遍历图片文件夹
for img_name in os.listdir(image_dir):
    if not img_name.lower().endswith(('.jpg', '.png', '.jpeg')):
        continue

    img_path = os.path.join(image_dir, img_name)

    # 读取图片
    img = cv2.imread(img_path)
    if img is None:
        print(f"无法读取图片: {img_name}")
        continue

    while True:

        # 创建临时图像副本用于显示绿色预览圆
        temp_img = img.copy()

        if 0 <= mouse_x < temp_img.shape[1] and 0 <= mouse_y < temp_img.shape[0]:
            cv2.circle(temp_img, (mouse_x, mouse_y), circle_radius, (0, 255, 0), 1)

        cv2.imshow("Paint Window", temp_img)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('+') or key == ord('='):
            circle_radius += 1
        elif key == ord('-') and circle_radius > 1:
            circle_radius -= 1
        elif key == 27:
            exit(0)
        elif key == ord('s'):
            output_img_path = os.path.join(output_dir, img_name)
            cv2.imwrite(output_img_path, img)
            print(f"已保存: {output_img_path}")
        elif key == ord('q'):
            break


cv2.destroyAllWindows()
# 完成
print("所有图像处理完成！")
